package tunnel

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
)

// ========================================
// 流式传输协议定义
// ========================================

const (
	// ChunkSize 每个分片的大小（1MB）
	ChunkSize = 1024 * 1024

	// MessageTypeResponseChunk 响应分片消息
	MessageTypeResponseChunk = 0x05

	// MessageTypeResponseEnd 响应结束消息
	MessageTypeResponseEnd = 0x06
)

// ChunkMessage 分片消息
type ChunkMessage struct {
	MessageType uint8
	TunnelID    string
	RequestID   string
	ChunkIndex  uint32 // 分片索引
	IsLastChunk bool   // 是否为最后一个分片
	ChunkLength uint32 // 本分片长度
	ChunkData   []byte // 分片数据
}

// NewChunkMessage 创建分片消息
func NewChunkMessage(tunnelID, requestID string, chunkIndex uint32, isLastChunk bool, data []byte) *ChunkMessage {
	return &ChunkMessage{
		MessageType: MessageTypeResponseChunk,
		TunnelID:    tunnelID,
		RequestID:   requestID,
		ChunkIndex:  chunkIndex,
		IsLastChunk: isLastChunk,
		ChunkLength: uint32(len(data)),
		ChunkData:   data,
	}
}

// Write 写入分片消息到Stream
func (m *ChunkMessage) Write(w io.Writer) error {
	// 写入消息头
	if err := binary.Write(w, binary.BigEndian, m.MessageType); err != nil {
		return err
	}

	// 写入TunnelID（36字节）
	tunnelIDBytes := []byte(m.TunnelID)
	if len(tunnelIDBytes) != 36 {
		return fmt.Errorf("invalid tunnelID length: %d", len(tunnelIDBytes))
	}
	if _, err := w.Write(tunnelIDBytes); err != nil {
		return err
	}

	// 写入RequestID（36字节）
	requestIDBytes := []byte(m.RequestID)
	if len(requestIDBytes) != 36 {
		return fmt.Errorf("invalid requestID length: %d", len(requestIDBytes))
	}
	if _, err := w.Write(requestIDBytes); err != nil {
		return err
	}

	// 写入分片索引
	if err := binary.Write(w, binary.BigEndian, m.ChunkIndex); err != nil {
		return err
	}

	// 写入是否最后一个分片
	var isLast uint8
	if m.IsLastChunk {
		isLast = 1
	}
	if err := binary.Write(w, binary.BigEndian, isLast); err != nil {
		return err
	}

	// 写入分片长度
	if err := binary.Write(w, binary.BigEndian, m.ChunkLength); err != nil {
		return err
	}

	// 写入分片数据
	if _, err := w.Write(m.ChunkData); err != nil {
		return err
	}

	return nil
}

// ReadChunkMessage 读取分片消息
func ReadChunkMessage(r io.Reader) (*ChunkMessage, error) {
	msg := &ChunkMessage{}

	// 读取消息类型
	if err := binary.Read(r, binary.BigEndian, &msg.MessageType); err != nil {
		return nil, err
	}

	// 读取TunnelID
	tunnelIDBytes := make([]byte, 36)
	if _, err := io.ReadFull(r, tunnelIDBytes); err != nil {
		return nil, err
	}
	msg.TunnelID = string(tunnelIDBytes)

	// 读取RequestID
	requestIDBytes := make([]byte, 36)
	if _, err := io.ReadFull(r, requestIDBytes); err != nil {
		return nil, err
	}
	msg.RequestID = string(requestIDBytes)

	// 读取分片索引
	if err := binary.Read(r, binary.BigEndian, &msg.ChunkIndex); err != nil {
		return nil, err
	}

	// 读取是否最后一个分片
	var isLast uint8
	if err := binary.Read(r, binary.BigEndian, &isLast); err != nil {
		return nil, err
	}
	msg.IsLastChunk = isLast == 1

	// 读取分片长度
	if err := binary.Read(r, binary.BigEndian, &msg.ChunkLength); err != nil {
		return nil, err
	}

	// 读取分片数据
	msg.ChunkData = make([]byte, msg.ChunkLength)
	if _, err := io.ReadFull(r, msg.ChunkData); err != nil {
		return nil, err
	}

	return msg, nil
}

// SendResponseStreaming 流式发送响应（支持大文件）
func SendResponseStreaming(w io.Writer, tunnelID, requestID string, reader io.Reader) error {
	buffer := make([]byte, ChunkSize)
	chunkIndex := uint32(0)
	totalBytes := int64(0)

	for {
		// 读取一个分片
		n, err := reader.Read(buffer)
		if n > 0 {
			isLast := err == io.EOF

			// 创建分片消息
			chunk := NewChunkMessage(tunnelID, requestID, chunkIndex, isLast, buffer[:n])

			// 发送分片
			if writeErr := chunk.Write(w); writeErr != nil {
				return fmt.Errorf("发送分片失败: %v", writeErr)
			}

			totalBytes += int64(n)
			chunkIndex++

			if chunkIndex%10 == 0 {
				log.Printf("📤 [STREAMING] 已发送 %d 个分片, 总大小: %.2f MB",
					chunkIndex, float64(totalBytes)/(1024*1024))
			}

			if isLast {
				log.Printf("✅ [STREAMING] 流式传输完成: %d 个分片, 总大小: %.2f MB",
					chunkIndex, float64(totalBytes)/(1024*1024))
				break
			}
		}

		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("读取数据失败: %v", err)
		}
	}

	return nil
}

// ReceiveResponseStreaming 流式接收响应（边接收边写入浏览器）
func ReceiveResponseStreaming(r io.Reader, w io.Writer, requestID string) error {
	chunkIndex := uint32(0)
	totalBytes := int64(0)

	for {
		// 读取分片消息
		chunk, err := ReadChunkMessage(r)
		if err != nil {
			return fmt.Errorf("读取分片失败: %v", err)
		}

		// 验证RequestID
		if chunk.RequestID != requestID {
			return fmt.Errorf("RequestID不匹配: 期望=%s, 实际=%s", requestID, chunk.RequestID)
		}

		// 验证分片索引
		if chunk.ChunkIndex != chunkIndex {
			return fmt.Errorf("分片索引不连续: 期望=%d, 实际=%d", chunkIndex, chunk.ChunkIndex)
		}

		// 写入分片数据到浏览器
		if _, err := w.Write(chunk.ChunkData); err != nil {
			return fmt.Errorf("写入分片数据失败: %v", err)
		}

		totalBytes += int64(chunk.ChunkLength)
		chunkIndex++

		if chunkIndex%10 == 0 {
			log.Printf("📥 [STREAMING] 已接收 %d 个分片, 总大小: %.2f MB",
				chunkIndex, float64(totalBytes)/(1024*1024))
		}

		// 检查是否为最后一个分片
		if chunk.IsLastChunk {
			log.Printf("✅ [STREAMING] 流式接收完成: %d 个分片, 总大小: %.2f MB",
				chunkIndex, float64(totalBytes)/(1024*1024))
			break
		}
	}

	return nil
}
